import os
import xml.etree.ElementTree as ET

VOC_CLASSES = (
    'aeroplane', 'bicycle', 'bird', 'boat',
    'bottle', 'bus', 'car', 'cat', 'chair',
    'cow', 'diningtable', 'dog', 'horse',
    'motorbike', 'person', 'pottedplant',
    'sheep', 'sofa', 'train', 'tvmonitor')


def parse_obj(xml_file):
    """parse a PASCAL VOC xml file"""
    tree = ET.parse(xml_file)
    objects = []
    for obj in tree.findall('object'):
        obj_struct = {}
        difficult = int(obj.find('difficult').text)
        if 1 == difficult:
            continue
        obj_struct['name'] = obj.find('name').text
        bbox = obj.find('bndbox')
        obj_struct['bbox'] = [int(float(bbox.find('xmin').text)),
                              int(float(bbox.find('ymin').text)),
                              int(float(bbox.find('xmax').text)),
                              int(float(bbox.find('ymax').text))]
        objects.append(obj_struct)

    return objects


def parse_xml2txt(xml_root, imgs_root, xml_files, txt_file_name, test_file_name):
    txt_file = open(txt_file_name, 'w')
    test_file = open(test_file_name, 'r')
    test_objs = test_file.readlines()
    test_imgs = [x.split('.')[0] for x in test_objs]

    for xml_file in xml_files:
        if xml_file.split('.')[0] in test_imgs:
            continue
        print(xml_file)
        img_name = xml_file.split('.')[0] + '.jpg'
        img_path = imgs_root + img_name
        if os.path.exists(img_path) is False:
            print('no such image', img_name)
            continue
        objects = parse_obj(xml_root + xml_file)
        if 0 == len(objects):
            print(xml_file + " has no object")
            continue

        txt_file.write(img_name)
        for obj in objects:
            cls_name = obj['name']
            bbox = obj['bbox']
            cls_idx = VOC_CLASSES.index(cls_name)
            txt_file.write(' ' + str(bbox[0]) + ' ' + str(bbox[1]) +
                           ' ' + str(bbox[2]) + ' ' + str(bbox[3]) +
                           ' ' + str(cls_idx))
        txt_file.write('\n')
    txt_file.close()
    test_file.close()


def xml2txt():
    annotations = '/media/tcl3/bokeh/lipeng/dataset/VOC2012/Annotations/'
    imgs_root = '/media/tcl3/bokeh/lipeng/dataset/VOC2012/JPEGImages/'
    xml_files = os.listdir(annotations)
    txt_file_name = '../voc2012.txt'
    test_file_name = '../voc2007test.txt'
    parse_xml2txt(annotations, imgs_root, xml_files, txt_file_name, test_file_name)


if __name__ == '__main__':
    xml2txt()
